#include "printf.h"
#include "pmm.h"
#include "vmm.h"
#include "interrupt.h"
#include "clint.h"
#include "riscv.h"
#include "mmio.h"
#include "sbi.h"
#include "scheduler.h"
#include "sync.h"       // 新增: 同步原语
#include <stdint.h>
#include <stddef.h>

// --------------------------
// UART寄存器定义
// --------------------------
#define UART0_ADDR 0x10000000UL
#define UART_RHR (UART0_ADDR + 0x0)
#define UART_THR (UART0_ADDR + 0x0)
#define UART_LSR (UART0_ADDR + 0x5)
#define UART_IER (UART0_ADDR + 0x1)
#define UART_IIR (UART0_ADDR + 0x2)

#define UART_LSR_RX_READY (1 << 0)
#define UART_LSR_TX_EMPTY (1 << 5)
#define UART_IER_RX_INT (1 << 0)

// --------------------------
// 延时函数
// --------------------------
void delay(int count) {
    for (int i = 0; i < count; i++) {
        for (int j = 0; j < 1000; j++) {
            __asm__ __volatile__("nop");
        }
    }
}

// --------------------------
// 定时器中断处理函数
// --------------------------
static volatile int timer_irq_count = 0;
static volatile int test_mode = 0;

void timer_handler(void) {
    timer_irq_count++;
    
    uint64_t next_time = get_time() + 1000000;
    sbi_set_timer(next_time);
    
    if (test_mode) {
        if (timer_irq_count <= 5) {
            printf("[中断测试] 第%d次时钟中断\n", timer_irq_count);
        }
        
        if (timer_irq_count == 5) {
            test_mode = 0;
            printf("[中断测试] 测试完成!\n\n");
        }
        return;
    }
    
    scheduler_tick();
    
    task_t* current = get_current_task();
    if (current) {
        current->elapsed_time += 1;
        
        if (timer_irq_count % 1000 == 0) {
            printf("\n[时钟中断] 第%d次触发(约%d秒), 当前任务: %s(PID=%d)\n",
                   timer_irq_count, timer_irq_count / 1000,
                   current->name, current->pid);
            print_tasks();
        }
        
        if (current->elapsed_time >= current->time_slice) {
            printf("[调度] 任务%s时间片耗尽,触发调度\n", current->name);
            schedule();
        }
    }
}

// --------------------------
// 中断功能测试
// --------------------------
void test_timer_interrupt(void) {
    printf("\n=== 中断功能测试 ===\n");
    printf("Testing timer interrupt...\n");
    
    timer_irq_count = 0;
    test_mode = 1;
    
    printf("等待5次中断...\n");
    
    while (test_mode) {
        wfi();
    }
    
    printf("Timer test completed: 5 interrupts\n");
    printf("=== 中断功能测试完成 ===\n\n");
}

// --------------------------
// UART功能
// --------------------------
static int uart_getc(void) {
    if (mmio_read32(UART_LSR) & UART_LSR_RX_READY) {
        return mmio_read32(UART_RHR) & 0xFF;
    } else {
        return -1;
    }
}

static void uart_init_irq(void) {
    mmio_write32(UART_IER, UART_IER_RX_INT);
}

static void uart_clear_irq(void) {
    mmio_read32(UART_IIR);
}

static volatile int uart_irq_count = 0;

void uart_handler(void) {
    uart_irq_count++;
    int c = uart_getc();
    if (c != -1) {
        printf("\n[UART中断] 第%d次触发，收到字符: '%c' (ASCII: %d)\n",
               uart_irq_count, c, c);
        uart_putc(c);
    }
    uart_clear_irq();
}

// --------------------------
// 虚拟地址访问测试
// --------------------------
void test_virtual_address_access() {
    void* phys_page = alloc_page();
    if (!phys_page) {
        printf("test_virtual_address_access: 分配物理页失败\n");
        return;
    }
    printf("\n=== 虚拟地址访问测试 ===\n");
    uint64_t pa = (uint64_t)phys_page;
    uint64_t va = KERNBASE + (pa - PHYSBASE);
    printf("物理页: PA = %p -> 虚拟地址: VA = %p\n", phys_page, (void*)va);
    printf("=== 虚拟地址访问测试结束 ===\n");
}

// ==========================================================
// 同步原语测试
// ==========================================================

// --------------------------
// 测试1: 自旋锁测试
// --------------------------
static spinlock_t test_spinlock;
static int shared_counter = 0;

void task_spinlock_test(void* arg) {
    int id = (int)(uintptr_t)arg;
    
    for (int i = 0; i < 3; i++) {
        spinlock_acquire(&test_spinlock);
        
        int old_value = shared_counter;
        printf("[自旋锁任务%d] 进入临界区, counter=%d\n", id, old_value);
        
        // 模拟临界区操作
        delay(100);
        shared_counter = old_value + 1;
        
        printf("[自旋锁任务%d] 离开临界区, counter=%d\n", id, shared_counter);
        
        spinlock_release(&test_spinlock);
        
        task_yield();
    }
    
    printf("[自旋锁任务%d] 完成\n", id);
}

// --------------------------
// 测试2: 条件变量 - 生产者消费者
// --------------------------
static spinlock_t buffer_lock;
static condvar_t not_empty;
static condvar_t not_full;
static int buffer[5];
static int buffer_count = 0;

void producer_condvar(void* arg) {
    int id = (int)(uintptr_t)arg;
    
    for (int i = 0; i < 5; i++) {
        spinlock_acquire(&buffer_lock);
        
        // 等待缓冲区不满
        while (buffer_count >= 5) {
            printf("[生产者%d] 缓冲区满,等待...\n", id);
            condvar_wait(&not_full, &buffer_lock);
        }
        
        // 生产数据
        int data = id * 10 + i;
        buffer[buffer_count++] = data;
        printf("[生产者%d] 生产数据: %d (缓冲区: %d/5)\n", 
               id, data, buffer_count);
        
        // 唤醒消费者
        condvar_signal(&not_empty);
        
        spinlock_release(&buffer_lock);
        
        delay(200);
    }
    
    printf("[生产者%d] 完成\n", id);
}

void consumer_condvar(void* arg) {
    int id = (int)(uintptr_t)arg;
    
    for (int i = 0; i < 5; i++) {
        spinlock_acquire(&buffer_lock);
        
        // 等待缓冲区不空
        while (buffer_count == 0) {
            printf("[消费者%d] 缓冲区空,等待...\n", id);
            condvar_wait(&not_empty, &buffer_lock);
        }
        
        // 消费数据
        int data = buffer[--buffer_count];
        printf("[消费者%d] 消费数据: %d (缓冲区: %d/5)\n", 
               id, data, buffer_count);
        
        // 唤醒生产者
        condvar_signal(&not_full);
        
        spinlock_release(&buffer_lock);
        
        delay(200);
    }
    
    printf("[消费者%d] 完成\n", id);
}

// --------------------------
// 测试3: 信号量 - 读者写者
// --------------------------
static semaphore_t mutex;          // 保护共享数据
static semaphore_t write_lock;     // 写者互斥
static int reader_count = 0;
static int shared_data = 0;

void reader_task(void* arg) {
    int id = (int)(uintptr_t)arg;
    
    for (int i = 0; i < 3; i++) {
        // 进入读者区
        semaphore_down(&mutex);
        reader_count++;
        if (reader_count == 1) {
            semaphore_down(&write_lock);  // 第一个读者锁定写者
        }
        semaphore_up(&mutex);
        
        // 读取数据
        printf("[读者%d] 读取数据: %d (读者数: %d)\n", 
               id, shared_data, reader_count);
        delay(150);
        
        // 离开读者区
        semaphore_down(&mutex);
        reader_count--;
        if (reader_count == 0) {
            semaphore_up(&write_lock);  // 最后一个读者释放写者
        }
        semaphore_up(&mutex);
        
        delay(200);
    }
    
    printf("[读者%d] 完成\n", id);
}

void writer_task(void* arg) {
    int id = (int)(uintptr_t)arg;
    
    for (int i = 0; i < 3; i++) {
        // 写者互斥
        semaphore_down(&write_lock);
        
        // 写入数据
        shared_data = id * 100 + i;
        printf("[写者%d] 写入数据: %d\n", id, shared_data);
        delay(150);
        
        semaphore_up(&write_lock);
        
        delay(300);
    }
    
    printf("[写者%d] 完成\n", id);
}

// --------------------------
// 同步原语综合测试
// --------------------------
void test_sync_primitives(void) {
    printf("\n=== 同步原语测试 ===\n");
    
    // 测试1: 自旋锁
    printf("\n[测试1] 自旋锁测试\n");
    spinlock_init(&test_spinlock, "test_lock");
    shared_counter = 0;
    
    create_task("SpinTask1", task_spinlock_test, (void*)1, 40);
    create_task("SpinTask2", task_spinlock_test, (void*)2, 40);
    
    delay(2000);
    printf("最终counter值: %d (期望6)\n", shared_counter);
    
    // 测试2: 条件变量(生产者-消费者)
    printf("\n[测试2] 条件变量 - 生产者消费者\n");
    spinlock_init(&buffer_lock, "buffer_lock");
    condvar_init(&not_empty, "not_empty");
    condvar_init(&not_full, "not_full");
    buffer_count = 0;
    
    create_task("Producer", producer_condvar, (void*)1, 50);
    create_task("Consumer", consumer_condvar, (void*)1, 50);
    
    delay(3000);
    
    // 测试3: 信号量(读者-写者)
    printf("\n[测试3] 信号量 - 读者写者\n");
    semaphore_init(&mutex, 1, "mutex");
    semaphore_init(&write_lock, 1, "write_lock");
    reader_count = 0;
    shared_data = 0;
    
    create_task("Reader1", reader_task, (void*)1, 50);
    create_task("Reader2", reader_task, (void*)2, 50);
    create_task("Writer1", writer_task, (void*)1, 50);
    
    delay(3000);
    
    printf("\n=== 同步原语测试完成 ===\n");
}

// --------------------------
// 主函数
// --------------------------
int main(void) {
    // 1-6. 前面的测试保持不变...
    printf("RISC-V 操作系统启动...\n");
    
    // ... (省略前面的初始化代码,保持不变) ...
    
    // 简化版主函数(只显示关键部分)
    pmm_init((uint64_t)_end, PHYSTOP);
    kvm_init();
    
    trap_init();
    int timer_reg = register_interrupt(IRQ_TIMER, timer_handler, IRQ_PRIORITY_3);
    if (timer_reg == 0) {
        enable_interrupt(IRQ_TIMER);
        uint64_t first_timer = get_time() + 1000000;
        sbi_set_timer(first_timer);
    }
    
    test_timer_interrupt();
    timer_irq_count = 0;
    
    int uart_reg = register_interrupt(IRQ_UART, uart_handler, IRQ_PRIORITY_5);
    if (uart_reg == 0) {
        uart_init_irq();
        enable_interrupt(IRQ_UART);
    }
    
    // 7. 调度器初始化
    scheduler_init();
    
    // 8. 同步原语测试(新增)
    test_sync_primitives();
    
    // 9. 内核主循环
    printf("\n内核运行中...\n");
    printf("功能特性:\n");
    printf("  - 轮转调度算法\n");
    printf("  - sleep/wakeup机制\n");
    printf("  - 自旋锁、条件变量、信号量\n");
    printf("  - 避免lost wakeup\n");
    printf("提示: 输入字符测试UART中断，按Ctrl+C退出QEMU\n\n");
    
    while (1) {
        wfi();
        
        static int loop_count = 0;
        if (++loop_count % 5000 == 0) {
            task_t* current = get_current_task();
            printf("[主循环] 循环计数: %d, 当前任务: %s\n",
                   loop_count, current ? current->name : "无");
        }
    }
    
    return 0;
}

